from paddlenlp import Taskflow
import argparse
import os
import json
import logging
import warnings
import re
import pdfplumber
from docx import opendocx, getdocumenttext
from win32com.client import Dispatch
from ecloud import CMSSEcloudOcrClient

warnings.filterwarnings("ignore")
logging.disable(logging.DEBUG)
logging.disable(logging.WARNING)

parser = argparse.ArgumentParser(description='Resume Parse')

parser.add_argument('--access_key', type=str, default='902ee66da8fc44f89735b500009f0a27')
parser.add_argument('--secret_key', type=str, default='3a0daa0843c74613bc1e1fe0551156ec')
parser.add_argument('--ecloud_url', type=str, default='https://api-wuxi-1.cmecloud.cn:8443')
parser.add_argument('--request_url', type=str, default='/api/ocr/v1/webimage')
parser.add_argument('--current_year', type=int, default=2024)
parser.add_argument('--schema', type=str, default='姓名,年龄,出生年份,毕业院校,学历')
parser.add_argument('--text', type=str, default=None)
parser.add_argument('--path', type=str, default=None)
args = parser.parse_args()

ie = Taskflow('information_extraction', schema=args.schema.split(','))

# 解析docx文件
def docx2str(path):

    docx_file = opendocx(path)
    text_list = getdocumenttext(docx_file)
    
    text_list_new = []
    for i, text in enumerate(text_list):
        if i != len(text_list) - 1:
            if text.find(text_list[i + 1]) != -1:
                continue
        if text[0] == '\t':
            text = text[1:]
        text = re.sub(r'\s+', '', text)
        if text not in text_list_new:
            text_list_new.append(text)
    text_str = ' '.join(text_list_new)
    # print(text_str)
    return text_str

# docx文件转doc文件
def doc2docx(path):

    if not os.path.isabs('.'):
        path = os.path.abspath(path)

    doc_file = Dispatch('Word.Application')
    doc = doc_file.Documents.Open(path)
    path_new = os.path.splitext(path)[0] + '.docx'
    doc.SaveAs(path_new, 12, False, "", True, "", False, False, False, False)
    doc.Close()
    doc_file.Quit()

    return path_new

# 解析pdf文件
def pdf2str(path):

    with pdfplumber.open(path) as pdf:
        text_list = []
        for page in pdf.pages:
            text_list.append(page.extract_text())
    text_str = ' '.join(text_list)
    # print(text_str)
    return text_str

# 解析图片
def img2str(path):

    try:
        ocr_client = CMSSEcloudOcrClient(args.access_key, args.secret_key, args.ecloud_url)
        response = ocr_client.request_ocr_service_file(requestpath=args.request_url, imagepath=path)
        words_info = json.loads(response.text)['body']['content']['prism_wordsInfo']
        text_str = ' '.join([word_info['word'] for word_info in words_info])
    except ValueError as e:
        print(e)
    
    # print(text_str)
    return text_str

# 信息抽取
def parse_resume(current_year, schema, text, path):

    if text is None:
        if os.path.exists(path):
            if path.endswith('txt'):
                with open(path, 'r', encoding='utf-8') as file:
                    text = ' '.join(file.readlines())
            elif path.endswith('docx'):
                text = docx2str(path)
            elif path.endswith('doc'):
                path_new = doc2docx(path)
                text = docx2str(path_new)
            elif path.endswith('pdf'):
                text = pdf2str(path)
            elif sum([path.endswith(ext) for ext in ['jpg', 'jpeg', 'png']]):
                text = img2str(path)
            else:
                return 'Invaild File Type:' + path
        else:
            return 'Invaild Path:' + path
    
    parse_res = ie(text)[0]
    res_dict = {}
    for parse_key in schema:
        if parse_key not in parse_res.keys():
            continue
        if parse_key == '出生年份':
            birth_yr = parse_res[parse_key][0]['text']
            if birth_yr.find('.') != -1:
                birth_yr = birth_yr.split('.')[0]
            age = current_year - int(birth_yr)
            if '年龄' not in parse_res.keys():
                res_dict['年龄'] = str(age)
            else:
                if parse_res[parse_key][0]['probability'] > parse_res['年龄'][0]['probability'] or len(res_dict['年龄']) > 2:
                    res_dict['年龄'] = str(age)
        else:
            res_dict[parse_key] = parse_res[parse_key][0]['text']

    res_str = ''
    for res_key in schema:
        if len(res_str) != 0:
            res_str += ' '
        if res_key in res_dict.keys():
            res_str += '1 ' + str(res_dict[res_key])
        else:
            res_str += '0'
    
    return res_str

res_str = parse_resume(args.current_year, args.schema.split(','), args.text, args.path)
# res_str = parse_resume(2024, args.schema.split(','), None, './resume_parser/test_data/1.docx')
print(res_str)
